#!/bin/bash
# *************************************
# * 功能: Shell脚本初始化Rocky环境
# * 作者: 王树森
# * 联系: wangshusen@sswang.com
# * 版本: 2025-06-11
# *************************************

# 错误处理：如果命令执行失败，脚本将终止
set -e

# 定义日志文件路径
LOG_FILE="/var/log/setup_script.log"

# 加载系统发行版信息
. /etc/os-release 

# 获取系统命令
[ "$ID" = "centos" -o "$ID" = "rocky" ] && os_cmd="yum" || os_cmd="apt"


# 日志记录函数
log() {
    local message="$1"
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    echo "[$timestamp] $message" | tee -a "$LOG_FILE"
}

# 定制yum源
customize_yum_sources() {
    if [ ! -d /etc/yum.repos.d/bak ]; then
        mkdir /etc/yum.repos.d/bak
        mv /etc/yum.repos.d/rocky* /etc/yum.repos.d/bak/
    fi
    source_list_file='/etc/yum.repos.d/base.repo'
    > "$source_list_file"
    for i in ali ustc; do
        [ "$i" == "ali" ] && domain="aliyun.com" || domain="ustc.edu.cn"
        [ "$i" == "ali" ] && os_name="rockylinux" || os_name="rocky"
        echo "# 定制 $i 源信息" >> "$source_list_file"
        for j in BaseOS AppStream; do
            cat >> "$source_list_file" <<-EOF
[$i-$j]
name=Rocky Linux \$releasever - $j - $domain
baseurl=https://mirrors.$domain/$os_name/\$releasever/$j/\$basearch/os/
gpgcheck=0
enabled=1
EOF
        done
    done
    ${os_cmd} install -y https://dl.fedoraproject.org/pub/epel/epel-release-latest-9.noarch.rpm
}

# 定制apt源
customize_apt_sources() {
    source_list_file='/etc/apt/sources.list'
    > "$source_list_file"
    for i in aliyun.com tuna.tsinghua.edu.cn; do
        log "定制 $i 源信息"
        echo "# 定制 $i 源信息" >> "$source_list_file"
        for j in noble noble-security noble-updates noble-backports; do
            for k in deb deb-src; do
                echo "$k https://mirrors.$i/ubuntu/ $j main restricted universe multiverse" >> "$source_list_file"
            done
        done
    done
}

# 定制源信息
customize_sources() {
    log "开始定制源信息..."
    if [ "${os_cmd}" == "yum" ]; then
        customize_yum_sources
    else
        customize_apt_sources
    fi
    log "源信息定制完成。"
}

# 更新软件源缓存
update_cache() {
    log "更新 软件源 缓存..."
    if [ "${os_cmd}" == "yum" ]; then
        ${os_cmd} makecache
    else
        rm -f /var/lib/dpkg/lock* /var/lib/apt/lists/lock
        ${os_cmd} update
    fi
    log "软件源 缓存更新完成。"
}

# 获取所有用户
get_all_user(){
    # 定制初始用户
    local userlist=""
    for i in $(ls /home); do
        userlist="$userlist $i"
    done
    echo "${userlist}"
}

# 获取所有bashrc文件
get_all_bashrc(){
    local bashrc_list="/root/.bashrc /etc/skel/.bashrc"
    userlist=$(get_all_user)
    for i in ${userlist}; do
        bashrc_list="${bashrc_list} /home/$i/.bashrc"
    done
    echo "${bashrc_list}"
}

# 定制yum源的提示符
customize_yum_prompt() {
    # 获取所有的bashrc文件
    local bashrc_list=$(get_all_bashrc)
    
    for i in ${bashrc_list}; do
    cat >> $i <<-eof
		# export PS1='\[\e[32m\][\u@\h \e[0m\e[35m\]\W\[\e[0m]\e[32m\]\]\\$\[\e[0m\] '
		export PS1='\033[36m\n\n当前位置：\033[31m\w \033[0m\n[\033[32m\u@\033[33m\h \033[36m\W \033[0m]\$ '
		eof
    done
}

# 定制apt的提示符
customize_apt_prompt() {
    # 获取所有的bashrc文件
    local bashrc_list=$(get_all_bashrc)

    # sed -i 's/^#force/force/; /033/s/w/W/; /34m/s/34m/35m/;' ${bashrc_list}
	for i in ${bashrc_list}; do
    cat >> $i <<-eof
		export PS1='\033[36m\n\n当前位置：\033[31m\w \033[0m\n[\033[32m\u@\033[33m\h \033[36m\W \033[0m]\$ '
		eof
    done
}

# 提示符信息定制
customize_prompt() {
    log "开始定制提示符信息..."
    if [ "${os_cmd}" == "yum" ]; then
        customize_yum_prompt
    else
        customize_apt_prompt
    fi
    log "提示符信息定制完成。"
}

# 安装基本软件工具
install_basic_tools() {
    log "开始安装基本软件工具..."
    
    # 安装通用软件
    local basesofts="vim net-tools bash-completion htop git tree psmisc rsyslog"
    local basemake="gcc make libxml2"
	rm -f /var/lib/dpkg/lock* /var/lib/apt/lists/lock
    ${os_cmd} install -y ${basesofts} ${basemake}
    
    # 安装定制的软件
    if [ "${os_cmd}" == "yum" ]; then
        ${os_cmd} install -y iputils httpd-tools
        ${os_cmd} install -y gcc-c++ glibc glibc-devel pcre pcre-devel openssl openssl-devel systemd-devel zlib-devel libxml2-devel
    else
        ${os_cmd} install -y iputils-ping apache2-utils
        ${os_cmd} install -y g++ build-essential libc6 libc6-dev libpcre3 libpcre3-dev libssl-dev libsystemd-dev zlib1g-dev libxml2-dev libxslt1-dev
    fi
    
    log "基本软件工具安装完成。"
}

# 定制pip基础环境
configure_pip_env(){
    ${os_cmd} install -y python3 python3-pip
    pip3 config set global.index-url https://mirrors.aliyun.com/pypi/simple/
    pip3 config set global.extra-index-url "https://pypi.tuna.tsinghua.edu.cn/simple https://pypi.doubanio.com/simple"
    pip3 config set install.trusted-host "mirrors.aliyun.com pypi.tuna.tsinghua.edu.cn"
}

# 定制 yum 系统的 python虚拟环境
configure_yum_python_venv() {
    pip3 install virtualenvwrapper
    cat >> "$HOME/.bashrc" <<-EOF
export WORKON_HOME=$HOME/.virtualenvs
export VIRTUALENVWRAPPER_PYTHON=/usr/bin/python3
source $vir_scripts
EOF
}

# 定制 apt 系统的 python虚拟环境
configure_apt_python_venv() {
    ${os_cmd} install -y python3-virtualenv python3-virtualenvwrapper
    cat >> "$HOME/.bashrc" <<-EOF
export WORKON_HOME=\$HOME/.virtualenvs
export VIRTUALENVWRAPPER_PYTHON=/usr/bin/python3
source $vir_scripts
EOF
}


# 配置 Python 虚拟环境
configure_python_venv() {
    log "开始配置 Python 虚拟环境..."
    
    # 准备环境
    configure_pip_env
    mkdir -p "$HOME/.virtualenvs"
    sed -i '/[Vv][Ii][Rr][Tt]/d' "$HOME/.bashrc"
    
    # 安装虚拟环境
    if [ "${os_cmd}" == "yum" ]; then
        vir_scripts='/usr/local/bin/virtualenvwrapper.sh'
        configure_yum_python_venv
    else
        vir_scripts='/usr/share/virtualenvwrapper/virtualenvwrapper.sh'
        configure_apt_python_venv
    fi
    
    # 虚拟环境收尾
    sed -i '/\-\-script/s#@"#@" >/dev/null 2>\&1#' $vir_scripts
    log "Python 虚拟环境配置完成。"
}

# 配置 Node.js 环境
configure_nodejs_env() {
    log "开始配置 Node.js 环境..."
    ${os_cmd}  install -y nodejs npm ruby
    npm config set registry https://registry.npmmirror.com
    npm config get registry
    log "Node.js 环境配置完成。"
}

# 配置 Golang 环境
configure_golang_env() {
    log "开始配置 Golang 环境..."
    ${os_cmd}  install -y golang
    sed -i -E '/GO(PROXY|SUMDB)/d' "$HOME/.bashrc"
    cat >> "$HOME/.bashrc" <<-EOF
export GOPROXY=https://goproxy.cn,direct
export GOSUMDB=sum.golang.google.cn
EOF
    log "Golang 环境配置完成。"
}

# 定制yum网卡信息
customize_yum_network_interface() {
    net_dir='/etc/NetworkManager/system-connections'
    rule_file='/etc/udev/rules.d/70.persistent-net.rules'
    mv $net_dir/ens160.nmconnection $net_dir/eth0.nmconnection
    sed -i 's/ens160/eth0/g' $net_dir/eth0.nmconnection
    mac_num=$(ip a show ens160 | awk '/ether/{print $2}')
    cat > ${rule_file} <<-EOF
SUBSYSTEM=="net", ACTION=="add", ATTR{address}=="$mac_num", ATTR{type}=="1", NAME="eth0"
EOF
    sed -i '/quiet"/s/quiet/quiet net.ifnames=0 biosdevname=0/' /etc/default/grub
    grub2-mkconfig -o /etc/grub2.cfg
}

# 定制apt网卡信息
customize_network_info_first() {
    sed -i 's/ens33/eth0/g' /etc/netplan/*.yaml
    sed -i '/LINUX="/s/="/="net.ifnames=0 biosdevname=0/' /etc/default/grub
    grub-mkconfig -o /boot/grub/grub.cfg
}

# 定制网卡信息
customize_network_interface() {
    log "开始定制网卡信息..."
    if [ "${os_cmd}" == "yum" ]; then
        customize_yum_network_interface
    else
        customize_network_info_first
    fi
    log "网卡信息定制完成。"
}

# 禁用 swap
disable_swap() {
    log "开始禁用 swap..."
    sed -i '/swap/d' /etc/fstab
    log "swap 禁用完成。"
}

# 基本安全加固
basic_security_hardening() {
    log "开始基本安全加固..."
    if [ "${os_cmd}" == "yum" ]; then
        sed -i 's/^SELINUX=.*/SELINUX=disabled/' /etc/selinux/config
        systemctl disable --now firewalld
    else
        systemctl disable --now ufw
    fi
    log "基本安全加固完成。"
}

# 设置时区
set_timezone() {
    log "开始配置时区..."
    timedatectl set-timezone Asia/Shanghai
    log "时区配置为 Asia/Shanghai。"
}

# 获取所有vimrc文件
get_all_vimrc(){
    local vimrc_list="/root/.vimrc /etc/skel/.vimrc"
    userlist=$(get_all_user)
    for i in ${userlist}; do
        vimrc_list="${vimrc_list} /home/$i/.vimrc"
    done
    echo "${vimrc_list}"
}

# 定制 vimrc
customize_vimrc() {
    log "开始定制 vimrc..."
	# 获取所有的vimrc文件
    local vimrc_list=$(get_all_vimrc)
    for i in ${vimrc_list}; do
        cat > ${i} <<-EOF
autocmd BufNewFile *.sh exec ":call ShellTitle()"
function ShellTitle()
    call append(0,"#!/bin/bash")
    call append(1,"# *************************************")
    call append(2,"# * 功能: Shell脚本模板")
    call append(3,"# * 作者: 王树森")
    call append(4,"# * 联系: wangshusen@sswang.com")
    call append(5,"# * 版本: ".strftime("%Y-%m-%d"))
    call append(6,"# *************************************")
endfunction
set tabstop=2
set shiftwidth=2
set expandtab
set number
set autoindent
set cursorline
set showmatch
syntax on
set background=dark
EOF
    done
    log "vimrc 定制完成。"
}

# 定制user的默认shell
customize_user_shell(){
    # 用户配置文件
    local user_conf="/etc/default/useradd"
    
    # 确保默认的shell是bash
    local user_shell=$(awk -F'=' '/^SHELL=/{print $2}' "${user_conf}")
    if [ "${user_shell}" != "/bin/bash" ]; then
      sed -i "/^SHELL=/s#sh#bash#" "${user_conf}"
    fi
}

# 定制ssh环境
customize_ssh(){
    # 定义配置文件路径
    config_file="/etc/ssh/sshd_config"

    # 安装ssh软件,为了避免安装异常版本冲突，这里增加非交互安装属性
    DEBIAN_FRONTEND=noninteractive ${os_cmd} install -y openssh-server
    
    if grep -q "PermitRootLogin" "$config_file"; then
        # 修改ssh配置
        sed -i '/PermitRootLogin/d' "${config_file}"
    fi  
    # 若不存在，则添加 PermitRootLogin yes
    echo "PermitRootLogin yes" >> "$config_file"	
}

# 基础功能环境定制
base_os_install(){
    log "开始基础功能环境定制..."
    customize_sources
    update_cache
    customize_prompt
    install_basic_tools
    customize_network_interface
    disable_swap
    basic_security_hardening
    set_timezone
    customize_vimrc
    customize_user_shell
    log "基础功能环境定制完成。"
}

# 基础应用环境定制
base_app_install(){
    log "开始基础应用环境定制..."
    configure_python_venv
    configure_nodejs_env
    configure_golang_env
    customize_ssh
    log "基础应用环境定制完成。"
}

# 主函数，按顺序调用各个功能函数
main() {
    log "开始系统环境定制..."
    base_os_install
    base_app_install
    log "所有操作已完成。"
    log "重启主机。"
    reboot
}

# 执行主函数
main
